from PySide6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QPushButton,
                               QFileDialog, QComboBox, QProgressBar, QLabel, QMessageBox, QLineEdit, QHBoxLayout,
                               QCheckBox)
from PySide6.QtCore import Qt, QThread, Signal
from ..translator.pdf_translator import PDFTranslator
from ..utils import ModelConfig
import os

class TranslationThread(QThread):
    update_progress = Signal(int)
    completed = Signal(bool, str)

    def __init__(self, file_path, target_lang, model_config, keep_layout):
        super().__init__()
        self.file_path = file_path
        self.target_lang = target_lang
        self.model_config = model_config
        self.keep_layout = keep_layout  # 新增布局保留参数

    def run(self):
        # 创建模型实例
        model = self.create_model()
        translator = PDFTranslator(model)
        output_path = os.path.splitext(self.file_path)[0] + '_translated.pdf'
        
        for progress in translator.translate_pdf(
            self.file_path, 
            target_language=self.target_lang,
            pages=None,
            output_file_path=output_path,
            keep_layout=self.keep_layout
        ):
            self.update_progress.emit(progress)
        
        self.completed.emit(True, output_path)

    def create_model(self):
        from ..model import OpenAIModel, GLMModel, DeepSeekModel  # 添加导入
        
        if self.model_config.model_type == "OpenAI":
            return OpenAIModel(self.model_config.api_key)
        elif self.model_config.model_type == "GLM":
            return GLMModel(self.model_config.api_key)
        elif self.model_config.model_type == "DeepSeek":
            # 传递模型名称和API密钥
            return DeepSeekModel(
                model="deepseek-chat",  # 具体模型名称
                api_key=self.model_config.api_key
            )
        else:
            raise ValueError("不支持的模型类型")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("PDF翻译工具")
        self.setGeometry(100, 100, 400, 200)
        self.init_ui()
        
    def init_ui(self):
        # 文件选择组件
        self.file_btn = QPushButton("选择PDF文件", self)
        self.file_btn.clicked.connect(self.select_file)
        
        # 语言选择
        self.lang_combo = QComboBox()
        self.lang_combo.addItems(["中文", "英文", "日语", "德语"])
        
        # 进度条
        self.progress = QProgressBar()
        self.progress.setAlignment(Qt.AlignCenter)
        
        # 新增模型配置组件
        self.model_combo = QComboBox()
        self.model_combo.addItems(["OpenAI", "GLM", "DeepSeek"])
        
        self.api_key_input = QLineEdit()
        self.api_key_input.setPlaceholderText("输入API密钥")
        
        # 新增操作按钮
        self.start_btn = QPushButton("开始翻译", self)
        self.start_btn.setEnabled(False)
        self.cancel_btn = QPushButton("取消", self)
        self.cancel_btn.setEnabled(False)
        
         # 新增布局保留复选框
        self.layout_checkbox = QCheckBox("保留原始版式布局")
        self.layout_checkbox.setChecked(True)  # 默认开启
        
        # 按钮布局
        button_layout = QHBoxLayout()
        button_layout.addWidget(self.start_btn)
        button_layout.addWidget(self.cancel_btn)
        
        
        
        # 布局
        layout = QVBoxLayout()
        layout.addWidget(QLabel("AI模型:"))
        layout.addWidget(self.model_combo)
        layout.addWidget(QLabel("API密钥:"))
        layout.addWidget(self.api_key_input)
        layout.addWidget(QLabel("选择PDF文件:"))
        layout.addWidget(self.file_btn)
        layout.addWidget(QLabel("目标语言:"))
        layout.addWidget(self.lang_combo)
        layout.addWidget(QLabel("翻译进度:"))
        layout.addWidget(self.progress)
        # 在原有布局中添加按钮布局
        layout.addLayout(button_layout)
        
        # 在布局中添加
        layout.addWidget(QLabel("高级选项:"))
        layout.addWidget(self.layout_checkbox)
        
        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)
        
        # 连接信号
        self.start_btn.clicked.connect(self.on_start_clicked)
        self.cancel_btn.clicked.connect(self.on_cancel_clicked)
    
    def select_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择PDF文件", "", "PDF文件 (*.pdf)")
        if file_path:
            self.file_path = file_path
            self.start_btn.setEnabled(True)  # 启用开始按钮
            
    def on_start_clicked(self):
        self.start_btn.setEnabled(False)
        self.cancel_btn.setEnabled(True)
        self.start_translation(self.file_path)
        
    def on_cancel_clicked(self):
        if self.thread and self.thread.isRunning():
            self.thread.terminate()
            self.progress.setValue(0)
            self.cancel_btn.setEnabled(False)
            QMessageBox.warning(self, "已取消", "翻译任务已中止")
    
    def start_translation(self, file_path):
        model_config = ModelConfig(
            model_type=self.model_combo.currentText(),
            api_key=self.api_key_input.text()
        )
        
        self.thread = TranslationThread(
            file_path,
            self.lang_combo.currentText(),
            model_config,
            keep_layout=self.layout_checkbox.isChecked()  # 传递复选框状态
        )
        self.thread.update_progress.connect(self.update_progress)
        self.thread.completed.connect(self.handle_result)
        self.thread.start()
    
    def update_progress(self, value):
        self.progress.setValue(value)
    
    def handle_result(self, success, message):
        if success:
            self.progress.setValue(100)
            QMessageBox.information(self, "完成", f"翻译文件已保存至：{message}")
        else:
            QMessageBox.critical(self, "错误", f"翻译失败：{message}")
